from scipy.spatial.transform import Rotation as sci_R
import numpy as np
from numpy.typing import NDArray
from typing import Tuple, Union

def euler_to_rotation(rot_xyz, degree=False):
    """ convert rotation in Euler angles to scipy's rotation
    :param rot_xyz: rotation in Euler angles, [r_x, r_y, r_z]
    :param degree: angle unit is rad or degree, default is False, use rad
    :return: rotation class in "scipy.spatial.transform.Rotation" type
    """
    return sci_R.from_euler('xyz', rot_xyz, degrees=degree)

def matrix_to_euler(rot_mat):
    """ convert rotation matrix to Euler angles
    :param rot_mat: rotation 3x3 matrix
    :return: rotation in Euler angles, [r_x, r_y, r_z]
    """
    return sci_R.from_matrix(rot_mat).as_euler('xyz')

def genertate_rt_matrix(translation: list, rotation: list, degree=False) -> NDArray:
    assert len(translation) == len(rotation) == 3
    mat = np.identity(4)
    mat[:3, :3] = euler_to_rotation(rotation, degree).as_matrix()
    mat[:3, 3] = np.array(translation)
    return mat

def matrix_to_rt(matrix: NDArray) -> Tuple[list, list]:
    assert matrix.shape == (4, 4)
    return matrix[:3, 3].tolist(), matrix_to_euler(matrix[:3, :3])

def pose_interpolate(timestamps: Union[list, NDArray],
                     timestamp1: Union[int, float], pose1: NDArray,
                     timestamp2: Union[int, float], pose2: NDArray) -> NDArray:
    '''
    timestamps: timestamps where we need to interpolate
    timestamp1, pose1: timestamp of pose1
    timestamp2, pose2: timestamp of pose2
    return pose interpolated at timestamp
    '''
    assert pose1.shape == (4, 4) and pose2.shape == (4, 4)
    if isinstance(timestamps, list):
        timestamps = np.array(timestamps)
    assert len(timestamps.shape) == 1
    # method1: timestamp1 <= timestamps <= timestamp2
    # from scipy.spatial.transform import Slerp
    # r12 = sci_R.from_matrix([pose1[:3, :3], pose2[:3, :3]])
    # slerp = Slerp([timestamp1, timestamp2], r12)
    # inter_rot = slerp(timestamps).as_matrix()
    scalers = (timestamps - timestamp1) / (timestamp2 - timestamp1)
    # method2
    rotvec12 = sci_R.from_matrix((pose2 @ np.linalg.inv(pose1))[:3, :3]).as_rotvec()
    rotvec1 = sci_R.from_matrix(pose1[:3, :3])
    inter_rot = (sci_R.from_rotvec(np.outer(scalers, rotvec12)) * rotvec1).as_matrix()
    # interpolate translation
    inter_trans = (np.outer(scalers, (pose2[:3, 3] - pose1[:3, 3])) + pose1[:3, 3])[..., np.newaxis]
    # compose poses
    inter_poses = np.concatenate((inter_rot, inter_trans), axis=-1)
    bottom = np.repeat(np.array([0, 0, 0, 1]).reshape(1, 1, 4), timestamps.shape[0], axis=0)
    inter_poses = np.concatenate((inter_poses, bottom), axis=1)
    return inter_poses
